# ray_tracer_utils.py

"""
射线追踪相关的几何计算工具。
包括射线与AABB、三角形以及BVH树的相交测试。
"""

import numpy as np
from .constants import EPSILON # 从新模块导入
from .bvh_tree import BVHNode   # 导入BVHNode定义

def ray_intersects_aabb(
    ray_origin_np: np.ndarray,
    inv_ray_direction_np: np.ndarray,
    ray_direction_signs_np: np.ndarray, # 0 for neg, 1 for pos/zero
    aabb_min_np: np.ndarray,
    aabb_max_np: np.ndarray
) -> bool:
    """检测射线是否与轴对齐包围盒 (AABB) 相交 (数值稳定版)。"""
    # 来自 voxelizer.py 的原始实现，这里不再重复注释
    tmin_val = -np.inf
    tmax_val = np.inf

    for i in range(3):
        if ray_direction_signs_np[i] == 1: 
            t_near = (aabb_min_np[i] - ray_origin_np[i]) * inv_ray_direction_np[i]
            t_far  = (aabb_max_np[i] - ray_origin_np[i]) * inv_ray_direction_np[i]
        else: 
            t_near = (aabb_max_np[i] - ray_origin_np[i]) * inv_ray_direction_np[i]
            t_far  = (aabb_min_np[i] - ray_origin_np[i]) * inv_ray_direction_np[i]
        
        tmin_val = max(tmin_val, t_near)
        tmax_val = min(tmax_val, t_far)

        if tmin_val > tmax_val:
            return False
            
    return tmax_val > EPSILON


def ray_intersects_triangle(
    ray_origin_np: np.ndarray, 
    ray_direction_np: np.ndarray, 
    triangle_np: tuple[np.ndarray, np.ndarray, np.ndarray]
) -> float | None: 
    """
    使用 Möller–Trumbore 算法检测射线是否与三角形相交，返回交点参数 t。
    """
    v0_np, v1_np, v2_np = triangle_np
    edge1_np = v1_np - v0_np
    edge2_np = v2_np - v0_np

    p_vec_np = np.cross(ray_direction_np, edge2_np)
    determinant = np.dot(edge1_np, p_vec_np)

    if -EPSILON < determinant < EPSILON: # abs(determinant) < EPSILON
        return None 

    inv_determinant = 1.0 / determinant
    t_vec_np = ray_origin_np - v0_np

    u = np.dot(t_vec_np, p_vec_np) * inv_determinant
    # 检查 u 是否在 [0, 1] 区间内 (包含边界)
    if u < (0.0 - EPSILON) or u > (1.0 + EPSILON) : 
        return None

    q_vec_np = np.cross(t_vec_np, edge1_np)
    v = np.dot(ray_direction_np, q_vec_np) * inv_determinant
    # 检查 v 是否在 [0, 1] 区间内 (包含边界) 且 u+v <= 1
    if v < (0.0 - EPSILON) or u + v > (1.0 + EPSILON) : 
        return None

    t = np.dot(edge2_np, q_vec_np) * inv_determinant

    if t > EPSILON: # 交点在射线正方向 (严格大于EPSILON，避免与起点重合或在其后)
        return t
    else:
        return None


def traverse_bvh_and_get_intersections(
    ray_origin_np: np.ndarray,
    ray_direction_np: np.ndarray,
    inv_ray_direction_np: np.ndarray,
    ray_direction_signs_np: np.ndarray,
    node: BVHNode | None
) -> list[float]:
    """
    递归遍历BVH树，收集射线与模型所有三角形的交点t值。
    """
    if node is None or node.bbox_min is None or node.bbox_max is None : 
        return []

    if not ray_intersects_aabb(ray_origin_np, inv_ray_direction_np, ray_direction_signs_np, node.bbox_min, node.bbox_max):
        return []

    intersections_t_values: list[float] = []
    if node.is_leaf:
        for triangle_np in node.triangles:
            t_value = ray_intersects_triangle(ray_origin_np, ray_direction_np, triangle_np)
            if t_value is not None:
                intersections_t_values.append(t_value)
    else:
        # 简单地测试左右子节点，未来可以根据t_min优化顺序
        if node.left_child:
            left_intersections = traverse_bvh_and_get_intersections(
                ray_origin_np, ray_direction_np, 
                inv_ray_direction_np, ray_direction_signs_np, 
                node.left_child
            )
            intersections_t_values.extend(left_intersections)

        if node.right_child:
            right_intersections = traverse_bvh_and_get_intersections(
                ray_origin_np, ray_direction_np,
                inv_ray_direction_np, ray_direction_signs_np,
                node.right_child
            )
            intersections_t_values.extend(right_intersections)
        
    return intersections_t_values

if __name__ == "__main__":
    print("--- 开始测试 ray_tracer_utils.py ---")

    # --- 测试 ray_intersects_triangle (返回 t 或 None) ---
    print("\n--- 测试 ray_intersects_triangle (返回 t | None) ---")
    test_triangle_np = (
        np.array([0.0, 0.0, 0.0], dtype=np.float64),
        np.array([1.0, 0.0, 0.0], dtype=np.float64),
        np.array([0.0, 1.0, 0.0], dtype=np.float64) # Triangle in XY plane at z=0
    )
    print(f"测试三角形 (NumPy): Vertices at z=0: v0=(0,0,0), v1=(1,0,0), v2=(0,1,0)")

    # (name, ray_origin, ray_direction, expected_t_is_not_None, approx_expected_t_if_intersects)
    test_cases_ray_tri = [
        ("射向内部 (z=1 to z=0)", [0.25, 0.25, 1.0], [0.0, 0.0, -1.0], True, 1.0),
        ("射向外部", [1.0, 1.0, 1.0], [0.0, 0.0, -1.0], False, None),
        ("平行平面", [0.25, 0.25, 1.0], [1.0, 0.0, 0.0], False, None),
        ("起点在内,方向远离", [0.25, 0.25, 0.0], [0.0, 0.0, 1.0], False, None), # t would be 0 or negative
        ("共面, determinant=0", [-0.1, -0.1, 0.0], [1.0, 1.0, 0.0], False, None),
        ("从背面射入 (z=-1 to z=0)", [0.25, 0.25, -1.0], [0.0, 0.0, 1.0], True, 1.0),
        ("起点近表面,方向离开", [0.25, 0.25, EPSILON / 2.0], [0.0, 0.0, 1.0], False, None), # t small negative
        ("精确击中顶点v0, 方向沿-Z", [0.0, 0.0, 1.0], [0.0, 0.0, -1.0], True, 1.0),
        ("精确击中边v0-v1中点, 方向沿-Z", [0.5, 0.0, 1.0], [0.0, 0.0, -1.0], True, 1.0),
    ]
    tri_passed_count = 0
    for name, origin_list, direction_list, expected_intersect, expected_t_val in test_cases_ray_tri:
        origin = np.array(origin_list, dtype=np.float64)
        direction = np.array(direction_list, dtype=np.float64)
        
        t_value = ray_intersects_triangle(origin, direction, test_triangle_np)
        intersects = (t_value is not None)
        
        print(f"  测试 '{name}': O={origin}, D={direction} -> Intersects: {intersects} (预期: {expected_intersect})", end="")
        if intersects:
            print(f", t={t_value:.4f}", end="")
            if expected_t_val is not None:
                 print(f" (预期 t~={expected_t_val:.4f})", end="")
        print()

        assert intersects == expected_intersect, f"Triangle Test '{name}' intersect-flag 失败"
        if intersects and expected_t_val is not None:
            assert np.isclose(t_value, expected_t_val, atol=EPSILON*10), \
                f"Triangle Test '{name}' t-value 失败: got {t_value}, expected {expected_t_val}"
        tri_passed_count +=1
    print(f"--- ray_intersects_triangle 测试通过 {tri_passed_count}/{len(test_cases_ray_tri)} ---")


    # --- 测试 ray_intersects_aabb ---
    print("\n--- 测试 ray_intersects_aabb ---")
    aabb_min = np.array([-1.0, -1.0, -1.0], dtype=np.float64)
    aabb_max = np.array([1.0, 1.0, 1.0], dtype=np.float64)
    print(f"测试AABB: min={aabb_min}, max={aabb_max}")

    test_cases_aabb = [
        ("中心射向+X", [0,0,0], [1,0,0], True),
        ("外部射向中心", [-2,0,0], [1,0,0], True),
        ("外部射失+X (Y方向射失)", [-2,2,0], [1,0,0], False),
        ("内部射向面外 (穿过角)", [0.5,0.5,0.5], [1,1,1], True),
        ("完全在后方", [2,0,0], [1,0,0], False), # t_max < 0
        ("起点在AABB内", [0,0,0], [-1,0,0], True),
        ("起点在AABB面外侧，平行于面", [-2,0,0], [0,1,0], False),
        ("精确穿过角点(-1,-1,-1)", [-2.1,-2.1,-2.1], [1,1,1], True),
        ("射线从角点旁边掠过(X轴方向)", [-1.1, -1.1, -1.1], [1,0,0], False),
        ("射线击中AABB远角", [-3,-3,-3], [1,1,0.5], True),
        ("射线从外部真正错过(斜向)", [-3,-3,-3], [1,1,0.4], False),
        ("方向分量为0 (沿X轴)", [-2,0,0], [1,0,0], True),
        ("方向分量为0 (沿Y轴，从外)", [0,-2,0], [0,1,0], True),
        ("方向分量为0 (沿Z轴，从内)", [0,0,0.5], [0,0,1], True),
        ("方向为0向量 (特殊情况)", [0,0,0], [0,0,0], False),
        ("射线起点在AABB的一个面上，向内射", [-1,0,0], [1,0,0], True), # t_min=0, t_max >0
        ("射线起点在AABB的一个面上，向外射", [-1,0,0], [-1,0,0], False),# t_max <0 (or t_min>t_max if origin on max face)
    ]
    aabb_passed_count = 0
    for name, origin_list, direction_list, expected in test_cases_aabb:
        origin = np.array(origin_list, dtype=np.float64)
        direction = np.array(direction_list, dtype=np.float64)
        
        inv_direction = np.zeros_like(direction)
        signs = np.zeros(3, dtype=int)

        if np.all(direction == 0):
            intersects = False 
        else:
            for i in range(3):
                if abs(direction[i]) < EPSILON: # Handle zero components in direction
                    inv_direction[i] = np.inf if direction[i] >= 0 else -np.inf
                else:
                    inv_direction[i] = 1.0 / direction[i]
                signs[i] = 1 if direction[i] >= 0 else 0 # 1 for positive/zero, 0 for negative
            intersects = ray_intersects_aabb(origin, inv_direction, signs, aabb_min, aabb_max)
            
        print(f"  测试 '{name}': O={origin}, D={direction} -> Intersects: {intersects} (预期: {expected})")
        assert intersects == expected, f"AABB测试 '{name}' 失败"
        aabb_passed_count += 1
    print(f"--- ray_intersects_aabb 测试通过 {aabb_passed_count}/{len(test_cases_aabb)} ---")

    # traverse_bvh_and_get_intersections 的测试会更复杂，
    # 通常与 voxelize_raycast(use_bvh=True) 的端到端测试一起进行，
    # 或者需要构造一个已知的小型BVH树和射线进行测试。
    # 暂时先依赖 voxelizer_core.py 中的端到端测试。

    print("--- ray_tracer_utils.py 测试结束 ---")